"""
启动playwright进程->启动浏览器->访问目标路径->飞书是否认证
    未授权认证/授权过期：走认证流程->手动授权->存储认证状态信息
    已授权过认证：访问目标路径->执行测试
"""
# 封装飞书认证的类：实现存储飞书认证状态，方便后续访问时可以不用登录
from playwright.sync_api import sync_playwright, Page, Browser, BrowserContext
import os.path
import time


class HesungUiAuto:
    # 初始化
    # 传入测试URL、飞书cookie存储的路径文件（有默认传参，可不传）、是否无头模式运行浏览器
    def __init__(self, target_url, auth_file_path="auth_state.json", headless=False):
        self.target_url = target_url            # 目标测试网站的URL
        self.auth_file_path = auth_file_path    # 飞书认证状态存储的路径，将cookie的json字段存储起来
        self.headless = headless                # 是否无头模式运行浏览器 是--则看不到运行窗口，否--运行浏览器窗口
        self.playwright = None                  # playwright driver 进程
        self.browser = None                     # Browser 类型的对象
        self.context = None
        self.page = None
    
    def __enter__(self):
        """支持上下文管理器模式使用"""
        self.start()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """退出时自动清理资源"""
        self.stop()

    # 启动浏览器 ，处理飞书认证逻辑
    def start(self):
        """启动playwright并处理认证"""
        # 启动playwright进程->启动浏览器->访问目标路径->飞书认证逻辑
        try:
            # 启动 playwright driver 进程
            self.playwright = sync_playwright().start()
            if not self.playwright:   # 如果playwright启动失败，则抛出异常
                raise Exception("无法启动Playwright")

            # 启动浏览器
            self.browser = self.playwright.chromium.launch(headless=self.headless)
            if not self.browser:   # 如果浏览器启动失败，则抛出异常
                raise Exception("无法启动浏览器")
            
            # 浏览器启动成功 则处理飞书认证逻辑
            self._handle_authentication()
            return self

        except Exception as e:
            print(f"启动过程中发生错误: {e}")
            self.stop()  # 确保资源被清理
            raise

    # 结束 关闭浏览器资源
    def stop(self):
        """停止并清理资源"""
        if self.browser:
            self.browser.close()
            self.browser = None
        if self.playwright:
            self.playwright.stop()
            self.playwright = None
        self.context = None
        self.page = None

    # 处理飞书认证逻辑
    def _handle_authentication(self):
        """处理认证逻辑，包括加载保存的认证状态或进行新的认证"""
        if not self.browser:
            raise Exception("浏览器未初始化")
            
        if os.path.exists(self.auth_file_path): # 使用已保存的认证状态
            
            print(f"使用已保存的认证状态: {self.auth_file_path}")
            self.context = self.browser.new_context(storage_state=self.auth_file_path)
            if not self.context:
                raise Exception("无法创建浏览器上下文")
                
            self.page = self.context.new_page()
            if not self.page:
                raise Exception("无法创建页面")
            
            # 访问目标页面
            print(f"访问目标页面: {self.target_url}")
            self.page.goto(self.target_url)
            
            # 如果在飞书认证页面，代表认证过期，需重新进行认证
            if self._is_feishu_auth_page():     # 判断是否在飞书认证页面
                print("认证状态已过期，需要重新认证")
                # 删除过期的认证状态
                os.remove(self.auth_file_path)
                # 执行新的认证流程
                self._perform_new_authentication()
        else:
            # 执行新的认证流程
            print("没有找到已保存的认证状态，需要进行认证")
            self.context = self.browser.new_context()
            if not self.context:
                raise Exception("无法创建浏览器上下文")
                
            self.page = self.context.new_page()
            if not self.page:
                raise Exception("无法创建页面")
                
            self._perform_new_authentication()

    # 飞书认证 手动认证授权
    def _perform_new_authentication(self):
        """执行新的认证流程"""
        if not self.page or not self.context:
            raise Exception("页面或上下文未初始化")
            
        # 访问目标页面，此时会重定向到飞书认证页面
        print(f"访问目标页面: {self.target_url}")
        self.page.goto(self.target_url)
        
        # 等待用户手动完成认证
        print("请在浏览器中完成飞书认证，认证完成后按回车...")
        input()
        
        # 等待认证完成并跳转
        print("等待认证完成并跳转...")
        self.page.wait_for_timeout(5000)
        
        # 等待重定向完成
        self._wait_for_redirect_complete()
        
        # 保存认证状态
        self.context.storage_state(path=self.auth_file_path)
        print(f"认证状态已保存到: {self.auth_file_path}")

    # 检查当前是否在飞书认证页面
    def _is_feishu_auth_page(self):
        return self.page and "accounts.feishu.cn" in self.page.url

    # 飞书认证完成后，重定向到目标url
    def _wait_for_redirect_complete(self, timeout=30):
        """等待重定向完成，最长等待timeout秒"""
        if not self.page:
            return False
        
        if self._is_feishu_auth_page():
            print("等待飞书认证完成并跳转...")
            for i in range(timeout):
                time.sleep(1)
                if not self._is_feishu_auth_page():
                    print(f"认证成功，已跳转到: {self.page.url}")
                    # 额外等待页面加载
                    self.page.wait_for_load_state("networkidle")
                    return True
            print(f"等待超时({timeout}秒)，认证可能失败")
            return False
        return True
    
    def get_page(self):
        """获取已认证的页面对象"""
        return self.page
    
    def get_context(self):
        """获取已认证的浏览器上下文"""
        return self.context
    
    def get_browser(self):
        """获取浏览器对象"""
        return self.browser
    
    def navigate_to(self, url):
        """导航到指定URL"""
        if self.page:
            self.page.goto(url)
            return self.page
        return None

    # 截取当前页面截图
    def take_screenshot(self, path):
        if self.page:
            return self.page.screenshot(path=path)
        return None

    # 开始跟踪
    def start_tracing(self):
        if self.context:
            self.context.tracing.start(snapshots=True, sources=True, screenshots=True)

    # 停止跟踪并保存zip文件
    def stop_tracing(self, path="trace.zip"):
        if self.context:
            self.context.tracing.stop(path=path)
            print(f"跟踪已保存到: {path}")


# 简单的使用示例
if __name__ == "__main__":
    try:
        # 示例：使用上下文管理器方式
        with HesungUiAuto("https://erp-dev.hesung.com/") as auth:
            # 获取已认证的页面
            page = auth.get_page()
            if page:
                print(f"当前页面标题: {page.title()}")
                
                # 开始跟踪
                auth.start_tracing()
                
                # 执行测试操作
                try:
                    # ---------------测试代码---------------
                    # 等待元素出现
                    menu_selector = 'ul.ant-menu-overflow>li'
                    page.wait_for_selector(menu_selector, timeout=10000)
                    
                    # 获取菜单项
                    menu_items = page.locator(menu_selector).all()
                    print(f"找到菜单项: {len(menu_items)} 个")
                    
                    # 输出菜单项文本
                    for i, item in enumerate(menu_items):
                        print(f"菜单 {i+1}: {item.inner_text()}")
                    
                    # 测试完成，停留几秒
                    print("测试完成，等待5秒...")
                    page.wait_for_timeout(5000)
                    
                except Exception as e:
                    print(f"测试过程中发生错误: {e}")
                    auth.take_screenshot("error.png")
                    print("错误截图已保存")
                
                # 停止跟踪
                auth.stop_tracing("demo_trace.zip")
    except Exception as e:
        print(f"程序执行过程中发生错误: {e}") 